import os
import numpy as np
import cPickle as pickle

class Plots(object):
    def __init__(self):
        self.plots = []
        self.tmp_dir = ''
        self.name = ''
    
    def add_plot( self, plot):
        plot.plots = self
        self.plots.append( plot )

    def plot_all( self ):
        for plot in self.plots:
            plot.plot()
    def read_all( self, std ):
        self.tmp_dir = std.tmp_dir
        self.name = std.name
        plot1 = PlotXYZ()
        plot1.read_conecyl( std.ccs[-1] )
        self.add_plot( plot1 )
        plot2 = PlotXYZ()
        plot2.read_study_cir( std )
        self.add_plot( plot2 )
        plot3 = PlotXYZ()
        plot3.read_study_mer( std )
        self.add_plot( plot3 )
    
    def save( self, path='' ):
        if path == '':
            path = os.path.join( self.tmp_dir, self.name + '.plots' )
        pfile = open( path, 'w' )
        pickle.dump( self, file=pfile, protocol=0 )
        pfile.close()

class PlotXYZ(object):
    
    def __init__(self):
        self.X = None
        self.Y = None
        self.Z = None
        self.plots = None
        
    def plot( self ):
        from mpl_toolkits.mplot3d import Axes3D
        from matplotlib import cm
        from matplotlib.ticker import LinearLocator, FormatStrFormatter
        import matplotlib.pyplot as plt
        fig = plt.figure()
        ax = fig.gca(projection='3d')
        surf = ax.plot_surface(self.X, self.Y, self.Z,
                               rstride=1, cstride=1, cmap=cm.jet,
                linewidth=0, antialiased=False)
        ax.set_zlim(-1.01, 1.01)

        ax.zaxis.set_major_locator(LinearLocator(10))
        ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

        fig.colorbar(surf, shrink=0.5, aspect=5)

        plt.show()

    def read_conecyl( self, cc ):
        step_name = cc.step2Name
        z_shape = len(cc.cross_sections)
        theta_shape = len(cc.meridians)
        X = np.array( np.zeros( (theta_shape, z_shape) ) )
        Y = np.array( np.zeros( (theta_shape, z_shape) ) )
        Z = np.array( np.zeros( (theta_shape, z_shape) ) )
        for z_index in range( z_shape ):
            for theta_index in range( theta_shape ):
                node = cc.cross_sections[z_index].nodes[theta_index]
                X[theta_index][z_index] = node.z
                Y[theta_index][z_index] = node.theta
                Z[theta_index][z_index] = node.fz[ step_name ]
        self.X = X
        self.Y = Y
        self.Z = Z

    def read_study_cir( self, study ):
        # circumferential displacements
        step_name = study.ccs[0].step1Name
        pl_shape = len( study.ccs )
        theta_shape = len( study.ccs[0].meridians )
        X = np.array( np.zeros( (theta_shape, pl_shape) ) )
        Y = np.array( np.zeros( (theta_shape, pl_shape) ) )
        Z = np.array( np.zeros( (theta_shape, pl_shape) ) )
        for pl_index in range( pl_shape ):
            cc = study.ccs[ pl_index ]
            pload = cc.ploadconf.ploads[0]
            cross_section = pload.cross_section
            for theta_index in range( theta_shape ):
                node = cross_section.nodes[ theta_index ]
                X[theta_index][pl_index] = pload.pltotal
                Y[theta_index][pl_index] = node.theta
                Z[theta_index][pl_index] = node.dr[ step_name ]
        self.X = X
        self.Y = Y
        self.Z = Z

    def read_study_mer( self, study ):
        # meridional displacements
        step_name = study.ccs[0].step1Name
        pl_shape = len( study.ccs )
        z_shape = len( study.ccs[0].cross_sections )
        X = np.array( np.zeros( (z_shape, pl_shape) ) )
        Y = np.array( np.zeros( (z_shape, pl_shape) ) )
        Z = np.array( np.zeros( (z_shape, pl_shape) ) )
        for pl_index in range( pl_shape ):
            cc = study.ccs[ pl_index ]
            pload = cc.ploadconf.ploads[0]
            meridian = pload.meridian
            for z_index in range( z_shape ):
                node = meridian.nodes[ z_index ]
                X[z_index][pl_index] = pload.pltotal
                Y[z_index][pl_index] = node.z
                Z[z_index][pl_index] = node.dr[ step_name ]
        self.X = X
        self.Y = Y
        self.Z = Z

